from sqlalchemy.ext.asyncio import create_async_engine, AsyncSession
from sqlalchemy.orm import sessionmaker, declarative_base
from sqlalchemy import text
from typing import Any, Optional, List

Base = declarative_base()


class AsyncDatabase:
    def __init__(self, db_url: str):
        self.engine = create_async_engine(
            db_url,
            echo=False,
            pool_size=20,
            max_overflow=10,
            pool_pre_ping=True
        )
        self.async_session = sessionmaker(
            bind=self.engine,
            class_=AsyncSession,
            expire_on_commit=False
        )

    async def execute_sql(self, sql: str, params: dict = None) -> Any:
        async with self.async_session() as session:
            try:
                result = await session.execute(text(sql), params or {})
                await session.commit()
                return result
            except Exception as e:
                await session.rollback()
                raise e

    async def fetch_all(self, sql: str, params: dict = None) -> List[dict]:
        result = await self.execute_sql(sql, params)
        return [dict(row) for row in result.mappings()]

    async def fetch_one(self, sql: str, params: dict = None) -> Optional[dict]:
        result = await self.fetch_all(sql, params)
        return result[0] if result else None

    async def insert(self, model: Base, data: dict) -> Any:
        async with self.async_session() as session:
            try:
                instance = model(**data)
                session.add(instance)
                await session.commit()
                await session.refresh(instance)
                return instance
            except Exception as e:
                await session.rollback()
                raise e

    async def update(self, model: Base, where: dict, data: dict) -> int:
        async with self.async_session() as session:
            try:
                stmt = text(
                    f"UPDATE {model.__tablename__} SET {', '.join([f'{k}=:{k}' for k in data.keys()])} WHERE {' AND '.join([f'{k}=:{k}_w' for k in where.keys()])}")
                params = {**data, **{f"{k}_w": v for k, v in where.items()}}
                result = await session.execute(stmt, params)
                await session.commit()
                return result.rowcount
            except Exception as e:
                await session.rollback()
                raise e

    async def delete(self, model: Base, where: dict) -> int:
        async with self.async_session() as session:
            try:
                stmt = text(
                    f"DELETE FROM {model.__tablename__} WHERE {' AND '.join([f'{k}=:{k}' for k in where.keys()])}")
                result = await session.execute(stmt, where)
                await session.commit()
                return result.rowcount
            except Exception as e:
                await session.rollback()
                raise e

    async def transaction(self):
        return self.async_session()

    async def fetch_pagination(self, sql: str, page: int = 1, per_page: int = 10, params: dict = None) -> dict:
        # 添加分页参数
        paginated_sql = f"{sql} LIMIT :limit OFFSET :offset"
        count_sql = f"SELECT COUNT(*) AS total FROM ({sql}) AS subquery"

        # 计算分页偏移
        offset = (page - 1) * per_page
        params = {**(params or {}), 'limit': per_page, 'offset': offset}

        # 执行查询
        data = await self.fetch_all(paginated_sql, params)
        total_result = await self.fetch_one(count_sql, params)

        return {
            'data': data,
            'pagination': {
                'current_page': page,
                'per_page': per_page,
                'total_pages': (total_result['total'] + per_page - 1) // per_page,
                'total_items': total_result['total']
            }
        }

    # 在AsyncDatabase类中添加以下方法
    async def get_total_count(self, sql: str, params: dict = None) -> int:
        count_sql = f"SELECT COUNT(*) AS total FROM ({sql}) AS subquery"
        result = await self.fetch_one(count_sql, params)
        return result['total'] if result else 0


if __name__ == '__main__':
    async def get_data():
        db = AsyncDatabase("mysql+aiomysql://root:lxy1qaz@101.34.74.27:3306/ct_decoration_prd")

        result = await db.fetch_pagination("SELECT * FROM internal_user", page=1, per_page=10)
        print(result)


    import asyncio

    asyncio.run(get_data())
